#!/bin/bash

# SSH Connectivity Test Script
# This script tests basic SSH connectivity to all VMs

set -e  # Exit on any error

# VM Configuration with updated IP addresses
VM1="192.168.1.206"  # Core Network
VM2="192.168.1.219"  # UPF
VM3="192.168.1.221"  # gNodeB
VM4="192.168.1.205"  # UE

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Logging function
log() {
    echo -e "${GREEN}[$(date +'%Y-%m-%d %H:%M:%S')] $1${NC}"
}

warn() {
    echo -e "${YELLOW}[$(date +'%Y-%m-%d %H:%M:%S')] WARNING: $1${NC}"
}

error() {
    echo -e "${RED}[$(date +'%Y-%m-%d %H:%M:%S')] ERROR: $1${NC}"
}

info() {
    echo -e "${BLUE}[$(date +'%Y-%m-%d %H:%M:%S')] INFO: $1${NC}"
}

# Function to test ping connectivity
test_ping() {
    local vm=$1
    local vm_name=$2
    
    log "Testing ping connectivity to $vm_name ($vm)..."
    
    if ping -c 3 -W 5 $vm >/dev/null 2>&1; then
        echo -e "✓ ${GREEN}Ping successful to $vm_name${NC}"
        return 0
    else
        echo -e "✗ ${RED}Ping failed to $vm_name${NC}"
        return 1
    fi
}

# Function to test SSH connectivity
test_ssh() {
    local vm=$1
    local user=$2
    local vm_name=$3
    
    log "Testing SSH connectivity to $vm_name ($user@$vm)..."
    
    if ssh -o ConnectTimeout=10 -o BatchMode=no $user@$vm "echo 'SSH connection successful'" 2>/dev/null; then
        echo -e "✓ ${GREEN}SSH successful to $vm_name${NC}"
        return 0
    else
        echo -e "✗ ${RED}SSH failed to $vm_name${NC}"
        return 1
    fi
}

# Function to get username for VM
get_username() {
    local vm=$1
    local vm_name=$2
    
    echo -n "Enter username for $vm_name ($vm) (default: ubuntu): "
    read username
    echo ${username:-ubuntu}
}

# Main test function
main() {
    echo
    echo "=========================================="
    echo "      SSH Connectivity Test"
    echo "=========================================="
    echo "Date: $(date)"
    echo "=========================================="
    echo
    
    log "Testing connectivity to all VMs:"
    log "  VM1 (Core): $VM1"
    log "  VM2 (UPF):  $VM2"
    log "  VM3 (gNB):  $VM3"
    log "  VM4 (UE):   $VM4"
    echo
    
    # Get usernames
    log "Please provide usernames for each VM:"
    USER1=$(get_username $VM1 "VM1 (Core)")
    USER2=$(get_username $VM2 "VM2 (UPF)")
    USER3=$(get_username $VM3 "VM3 (gNB)")
    USER4=$(get_username $VM4 "VM4 (UE)")
    
    echo
    log "Using usernames:"
    log "  VM1 ($VM1): $USER1"
    log "  VM2 ($VM2): $USER2"
    log "  VM3 ($VM3): $USER3"
    log "  VM4 ($VM4): $USER4"
    echo
    
    # Test ping connectivity
    log "Testing ping connectivity..."
    ping_results=0
    
    test_ping $VM1 "VM1 (Core)" && ((ping_results++))
    test_ping $VM2 "VM2 (UPF)" && ((ping_results++))
    test_ping $VM3 "VM3 (gNB)" && ((ping_results++))
    test_ping $VM4 "VM4 (UE)" && ((ping_results++))
    
    echo
    log "Ping Results: $ping_results/4 successful"
    echo
    
    # Test SSH connectivity
    log "Testing SSH connectivity..."
    ssh_results=0
    
    test_ssh $VM1 $USER1 "VM1 (Core)" && ((ssh_results++))
    test_ssh $VM2 $USER2 "VM2 (UPF)" && ((ssh_results++))
    test_ssh $VM3 $USER3 "VM3 (gNB)" && ((ssh_results++))
    test_ssh $VM4 $USER4 "VM4 (UE)" && ((ssh_results++))
    
    echo
    log "SSH Results: $ssh_results/4 successful"
    echo
    
    # Summary
    echo "=========================================="
    echo "           Test Summary"
    echo "=========================================="
    echo "Ping Tests: $ping_results/4 passed"
    echo "SSH Tests:  $ssh_results/4 passed"
    echo "=========================================="
    
    if [ $ping_results -eq 4 ] && [ $ssh_results -eq 4 ]; then
        log "✓ All connectivity tests passed!"
        log "You can now run the SSH key setup script."
        echo
        log "Next step: ./scripts/setup-ssh-keys.sh"
    else
        error "Some connectivity tests failed!"
        echo
        error "Please check:"
        error "  1. All VMs are running and accessible"
        error "  2. Network connectivity between jump host and VMs"
        error "  3. SSH service is running on all VMs"
        error "  4. Usernames are correct"
        error "  5. Firewall rules allow SSH connections"
    fi
}

# Handle command line arguments
case "${1:-}" in
    "--help"|"-h")
        echo "Usage: $0 [OPTIONS]"
        echo
        echo "Options:"
        echo "  -h, --help          Show this help message"
        echo
        echo "This script will:"
        echo "  1. Test ping connectivity to all VMs"
        echo "  2. Test SSH connectivity to all VMs"
        echo "  3. Provide a summary of connectivity status"
        echo
        echo "Prerequisites:"
        echo "  - All VMs must be running"
        echo "  - Network connectivity must be established"
        echo "  - SSH service must be running on all VMs"
        ;;
    *)
        main
        ;;
esac 